from abc import ABC, abstractmethod
import numpy as np
from typing import Union


class AbstractBayesClassifier(ABC):

    def __init__(self):
        self._is_fitted = False

    def fit(self, train_data: np.array, train_target: np.array) -> None:
        """
        该方法拟合训练数据
        @param train_data: 训练集数据
        @param train_target: 训练集标签
        @return: None
        """
        self._implemented_fit(train_data, train_target)
        self._is_fitted = True

    @abstractmethod
    def _implemented_fit(self, train_data: np.array, train_target: np.array) -> None:
        """
        拟合数据抽象方法，子类需实现该方法
        @param train_data: 训练集数据
        @param train_target: 训练集标签
        @return: None
        """
        pass

    def predict(self, data: Union[list, np.array]) -> np.array:
        """
        调用模型产生预测
        @param data: 需要预测的数据
        @return: 预测的结果
        """
        if not self._is_fitted:
            raise ValueError('Cannot predict before fitting the model.')
        return self._implemented_predict(data)

    @abstractmethod
    def _implemented_predict(self, data: Union[list, np.array]) -> np.array:
        """
        预测方法，子类需实现
        @param data: 需要预测的数据
        @return: 预测的结果
        """
        pass


class NaiveBayesClassifier(AbstractBayesClassifier):
    def __init__(self):
        super().__init__()
        self.bayes_dict = {}

    def _implemented_fit(self, train_data: np.array, train_target: np.array) -> None:
        # 计算先验概率
        self.bayes_dict['prior_probability'] = self.calculate_proportions(train_target.reshape(-1))
        # 计算条件概率
        self.bayes_dict['conditional_probability'] = {}
        for feature_index in range(len(train_data[0])):
            self.bayes_dict['conditional_probability'][feature_index] = ProbabilityCalculator(
                train_data[:, feature_index], train_target)
        print('Data fitting successful!')

    def _implemented_predict(self, data: np.array) -> np.array:
        result = []
        # 遍历输入数据，对数据点逐一预测
        for data_point in data:
            # 获取所有的类别
            label_names = list(self.bayes_dict['prior_probability'].keys())
            single_result = []
            # 遍历各类别，计算联合概率分布
            for label_name in label_names:
                # 获取先验概率
                jpd = [self.bayes_dict['prior_probability'][label_name]]
                # 遍历获取每一个特征的条件概率
                for feature_idx in self.bayes_dict['conditional_probability']:
                    conditional_probability = self.bayes_dict['conditional_probability'][feature_idx].get_probability(
                        data_point[feature_idx], label_name)
                    jpd.append(conditional_probability)
                jpd = np.array(jpd)
                jpd = np.prod(jpd)
                single_result.append(jpd)

            result.append(single_result)
        # print(result)
        # 取最大值索引返回
        return np.argmax(result, axis=-1)

    @staticmethod
    def calculate_proportions(array: np.array) -> dict:
        """
        此方法用于计算先验概率
        @param array: 类别标签
        @return: 字典，各标签（key）及其频率（value）
        """
        # 初始化一个空字典来存储结果
        result = {}

        # 计算数组的长度
        total_length = len(array)

        # 对数组中的每个元素进行计数
        for num in array:
            if num in result:
                result[num] += 1
            else:
                result[num] = 1

        # 计算每个元素的占比
        for key in result:
            result[key] = result[key] / total_length

        return result


class ProbabilityCalculator:
    """
    概率计算器类
    """

    def __init__(self, feature: np.array, target: np.array) -> None:
        """
        初始化函数，根据特征的类型（连续或离散）来处理概率模型
        """
        self.prob_model = None
        # 确保特征的类型一致
        assert len(set([type(x) for x in feature])) == 1, ValueError
        if type(feature[0]) == float:
            self.type = 'continuous'
            # 如果特征是连续的，处理连续概率
            self._process_continuous_probability(feature, target)
        elif type(feature[0]) == int:
            self.type = 'discrete'
            # 如果特征是离散的，处理离散概率
            self._process_discrete_probability(feature, target)
        else:
            raise NotImplementedError

    def _process_discrete_probability(self, feature: np.array, target: np.array) -> None:
        """
        处理离散概率，计算特征和标签的联合频率
        """
        feature = feature.reshape(-1)
        target = target.reshape(-1)

        joint_freq = np.zeros((np.max(feature) + 1, np.max(target) + 1))
        for f, l in zip(feature, target):
            joint_freq[f, l] += 1

        self.prob_model = joint_freq

    def _process_continuous_probability(self, feature: np.array, target: np.array) -> None:
        """
        处理连续概率，对每个标签，计算特征的均值和标准差
        """
        feature = feature.reshape(-1)
        target = target.reshape(-1)

        grouped_data = {}

        for d, l in zip(feature, target):
            if l in grouped_data:
                grouped_data[l].append(d)
            else:
                grouped_data[l] = [d]

        self.prob_model = {}
        for label in grouped_data:
            mu = np.mean(grouped_data[label])
            sigma = np.std(grouped_data[label], ddof=1)

            self.prob_model[label] = {}
            self.prob_model[label]['mu'] = mu
            self.prob_model[label]['sigma'] = sigma

    @staticmethod
    def _calculate_gaussian_probability(x: float, mu: float, sigma: float) -> float:
        """
        计算高斯概率
        """
        sigma_squared = np.power(sigma, 2)
        exponent = np.exp(-(np.power(x - mu, 2)) / (2 * sigma_squared))
        p = (1 / (np.sqrt(2 * np.pi) * sigma)) * exponent
        return p

    def _calculate_discrete_probability(self, feature_value: np.array, target: np.array, _lambda=1) -> float:
        """
        计算离散概率
        """
        freq = self.prob_model
        max_feature_value_index = len(freq) - 1
        while feature_value > max_feature_value_index:
            row = np.zeros_like(freq[0])
            freq = np.vstack((freq, row))
            max_feature_value_index = len(freq) - 1
        freq += _lambda
        s = sum(freq[:, target])
        freq = freq[:, target] / s
        return freq[feature_value]

    def get_probability(self, feature: np.array, target: np.array) -> float:
        """
        获取特征值在目标类别下的概率
        """
        if self.type == 'discrete':
            return self._calculate_discrete_probability(feature, target)
        elif self.type == 'continuous':
            mu = self.prob_model[target]['mu']
            sigma = self.prob_model[target]['sigma']
            return self._calculate_gaussian_probability(feature, mu, sigma)
        else:
            raise NotImplementedError
